import win32gui
import win32ui
import win32con
import numpy as np
import cv2
import time
import threading
from enum import Enum

class DrawColor(Enum):
    """绘制颜色"""
    RED = 0x0000FF     # BGR格式
    GREEN = 0x00FF00
    BLUE = 0xFF0000
    YELLOW = 0x00FFFF
    PURPLE = 0xFF00FF
    WHITE = 0xFFFFFF
    BLACK = 0x000000

class WindowDraw:
    # 存储永久显示的图形
    _permanent_shapes = {}
    # 控制绘制线程
    _draw_thread = None
    _running = False

    @classmethod
    def _draw_permanent_shapes(cls):
        """持续绘制永久图形的线程函数"""
        while cls._running:
            for hwnd, shapes in cls._permanent_shapes.items():
                if not shapes:
                    continue
                
                try:
                    # 获取DC
                    hdc = win32gui.GetDC(hwnd)
                    mfc_dc = win32ui.CreateDCFromHandle(hdc)
                    
                    for shape in shapes:
                        if shape['type'] == 'image':
                            # 创建内存DC
                            mem_dc = mfc_dc.CreateCompatibleDC()
                            
                            # 创建位图
                            bitmap = win32ui.CreateBitmap()
                            bitmap.CreateCompatibleBitmap(mfc_dc, shape['width'], shape['height'])
                            mem_dc.SelectObject(bitmap)
                            
                            # 将numpy数组转换为位图
                            image = shape['image']
                            height, width = image.shape[:2]
                            
                            # 使用BitBlt直接绘制
                            mem_dc.BitBlt(
                                (0, 0),
                                (width, height),
                                mfc_dc,
                                (shape['x'], shape['y']),
                                win32con.SRCCOPY
                            )
                            
                            # 使用TransparentBlt实现透明效果
                            win32gui.TransparentBlt(
                                hdc,                    # 目标DC
                                shape['x'],             # 目标x
                                shape['y'],             # 目标y
                                shape['width'],         # 目标宽度
                                shape['height'],        # 目标高度
                                mem_dc.GetHandleOutput(), # 源DC
                                0, 0,                   # 源x,y
                                shape['width'],         # 源宽度
                                shape['height'],        # 源高度
                                0xFF00FF                # 透明色（洋红色）
                            )
                            
                            # 清理资源
                            mem_dc.DeleteDC()
                            win32gui.DeleteObject(bitmap.GetHandle())
                        
                        elif shape['type'] == 'rect':
                            # 原有的矩形绘制代码
                            pen = win32gui.CreatePen(win32con.PS_SOLID, shape['thickness'], shape['color'].value)
                            old_pen = win32gui.SelectObject(hdc, pen)
                            win32gui.MoveToEx(hdc, shape['x'], shape['y'])
                            win32gui.LineTo(hdc, shape['x'] + shape['width'], shape['y'])
                            win32gui.LineTo(hdc, shape['x'] + shape['width'], shape['y'] + shape['height'])
                            win32gui.LineTo(hdc, shape['x'], shape['y'] + shape['height'])
                            win32gui.LineTo(hdc, shape['x'], shape['y'])
                            win32gui.SelectObject(hdc, old_pen)
                            win32gui.DeleteObject(pen)
                    
                    # 清理资源
                    mfc_dc.DeleteDC()
                    win32gui.ReleaseDC(hwnd, hdc)
                    
                except Exception as e:
                    print(f"绘制图形失败: {e}")
                
                time.sleep(0.016)  # 约60fps的刷新率

    @classmethod
    def _ensure_thread_running(cls):
        """确保绘制线程在运行"""
        if not cls._running:
            cls._running = True
            cls._draw_thread = threading.Thread(target=cls._draw_permanent_shapes, daemon=True)
            cls._draw_thread.start()

    @classmethod
    def add_permanent_rect(cls, hwnd: int, x: int, y: int, width: int, height: int,
                          color: DrawColor = DrawColor.RED, thickness: int = 2):
        """添加永久显示的矩形"""
        if hwnd not in cls._permanent_shapes:
            cls._permanent_shapes[hwnd] = []
        
        cls._permanent_shapes[hwnd].append({
            'type': 'rect',
            'x': x,
            'y': y,
            'width': width,
            'height': height,
            'color': color,
            'thickness': thickness
        })
        
        cls._ensure_thread_running()

    @classmethod
    def add_permanent_text(cls, hwnd: int, text: str, x: int, y: int,
                          color: DrawColor = DrawColor.RED, font_size: int = 20):
        """添加永久显示的文字"""
        if hwnd not in cls._permanent_shapes:
            cls._permanent_shapes[hwnd] = []
        
        cls._permanent_shapes[hwnd].append({
            'type': 'text',
            'text': text,
            'x': x,
            'y': y,
            'color': color,
            'font_size': font_size
        })
        
        cls._ensure_thread_running()

    @classmethod
    def clear_permanent_shapes(cls, hwnd: int):
        """清除指定窗口的所有永久图形"""
        if hwnd in cls._permanent_shapes:
            cls._permanent_shapes[hwnd] = []
            win32gui.InvalidateRect(hwnd, None, True)

    @classmethod
    def stop(cls):
        """停止所有绘制"""
        cls._running = False
        if cls._draw_thread:
            cls._draw_thread.join()
        cls._permanent_shapes.clear()

    @staticmethod
    def draw_rect(hwnd: int, 
                  x: int, 
                  y: int, 
                  width: int, 
                  height: int, 
                  color: DrawColor = DrawColor.RED,
                  thickness: int = 2,
                  duration_ms: int = 500,
                  auto_clear: bool = True):
        """临时绘制矩形"""
        def draw():
            try:
                dc = win32gui.GetDC(hwnd)
                pen = win32gui.CreatePen(win32con.PS_SOLID, thickness, color.value)
                old_pen = win32gui.SelectObject(dc, pen)
                
                start_time = time.time()
                while time.time() - start_time < duration_ms/1000:
                    win32gui.Rectangle(dc, x, y, x + width, y + height)
                    time.sleep(0.016)  # 约60FPS
                
                win32gui.SelectObject(dc, old_pen)
                win32gui.DeleteObject(pen)
                win32gui.ReleaseDC(hwnd, dc)
                
                if auto_clear:
                    win32gui.InvalidateRect(hwnd, (x-thickness, y-thickness, 
                                                 x + width+thickness, y + height+thickness), True)
            except Exception as e:
                print(f"临时绘制矩形失败: {e}")
        
        # 在新线程中绘制
        threading.Thread(target=draw, daemon=True).start()

    @staticmethod
    def draw_text(hwnd: int,
                  text: str,
                  x: int,
                  y: int,
                  color: DrawColor = DrawColor.RED,
                  font_size: int = 20,
                  duration_ms: int = 500,
                  auto_clear: bool = True):
        """临时绘制文字"""
        def draw():
            try:
                dc = win32gui.GetDC(hwnd)
                win32gui.SetTextColor(dc, color.value)
                win32gui.SetBkMode(dc, win32con.TRANSPARENT)
                
                start_time = time.time()
                while time.time() - start_time < duration_ms/1000:
                    win32gui.DrawText(dc, text, -1, 
                                    (x, y, x + 500, y + font_size),
                                    win32con.DT_LEFT | win32con.DT_NOCLIP)
                    time.sleep(0.016)  # 约60FPS
                
                win32gui.ReleaseDC(hwnd, dc)
                
                if auto_clear:
                    win32gui.InvalidateRect(hwnd, (x, y, x + 500, y + font_size), True)
            except Exception as e:
                print(f"临时绘制文字失败: {e}")
        
        # 在新线程中绘制
        threading.Thread(target=draw, daemon=True).start()

    @staticmethod
    def draw_rect_once(hwnd: int, 
                       x: int, 
                       y: int, 
                       width: int, 
                       height: int, 
                       color: DrawColor = DrawColor.RED,
                       thickness: int = 2):
        """绘制一次性矩形（不会自动消失）"""
        try:
            dc = win32gui.GetDC(hwnd)
            pen = win32gui.CreatePen(win32con.PS_SOLID, thickness, color.value)
            old_pen = win32gui.SelectObject(dc, pen)
            
            win32gui.Rectangle(dc, x, y, x + width, y + height)
            
            win32gui.SelectObject(dc, old_pen)
            win32gui.DeleteObject(pen)
            win32gui.ReleaseDC(hwnd, dc)
        except Exception as e:
            print(f"绘制一次性矩形失败: {e}")

    @classmethod
    def add_permanent_image(cls, hwnd: int, image: np.ndarray, x: int, y: int, width: int, height: int):
        """添加永久显示的图像"""
        if hwnd not in cls._permanent_shapes:
            cls._permanent_shapes[hwnd] = []
        
        cls._permanent_shapes[hwnd].append({
            'type': 'image',
            'image': image,
            'x': x,
            'y': y,
            'width': width,
            'height': height
        })
        
        cls._ensure_thread_running()